import numpy as np
import torch
import time
import math
torch.set_printoptions(8)

def gelu(x):
    """
        Task: Use the torch API to implement the approximate calculation formula of the `GELU`
        activation function. The formula is as follows (you need to paste it into the latex
        online conversion website)
        Website: https://www.latexlive.com/
        Formula: \frac{1}{2} x\left[1+\tanh \left(\sqrt{\frac{2}{\pi}}\left(x+0.044715 x^{3}\right)\right)\right]
        
        Input: Tensor
        Output: Tensor
    """
    const_temp_val = math.sqrt(2/ math.pi)

    return 0.5 * x * (1 + torch.tanh(const_temp_val * (x + 0.044715 *  x**3)))


def softmax(x, dim=-1):
    """
        Task: Use torch API to implement `softmax` function, search the specific formula by yourself
        Input: Tensor
        Output: Tensor
    """
    max_vals = torch.max(x, dim=dim, keepdim=True).values
    x_exp = torch.exp(x - max_vals)
    sum_exp = torch.sum(x_exp, dim=dim, keepdim=True)
    return x_exp / sum_exp

def layer_norm(x, g_b, eps:float = 1e-5):
    """
        Task: Use torch API to implement `layernorm` function, search `layernorm` by yourself
        Input: 
            x: Tensor
            g_b: dictionary that load from gpt2 weight. g-gamma and b-bias are the keys
        Output: Tensor
    """
    g, b = torch.Tensor(g_b['g']), torch.Tensor(g_b['b'])
    mean = x.mean(dim=-1, keepdim=True)
    var  = x.var(dim=-1, keepdim=True, unbiased=False)

    out = (x - mean) * torch.rsqrt(var + eps)

    return g * out +b

def linear(x, w_b):  # [m, in], [in, out], [out] -> [m, out]
    """
        Task: implement linear layer 
        Input: 
            x: Tensor
            w_b: dictionary that load from gpt2 weight. w-weight and b-bias are the keys
        Output: Tensor
    """
    w, b = w_b['w'], w_b['b']
    
    if isinstance(w, np.ndarray):
        w = torch.from_numpy(w)

    out = torch.matmul(x, w)

    if b is not None:
        out += b

    return out
    

def ffn(x, mlp):  # [n_seq, n_embd] -> [n_seq, n_embd]
    """
        Task: use `gelu` `linear` to implement ffn
        Notes: x --linear--> --gelu--> --linear--> output
        Input: 
            x: Tensor
            mlp: dictionary that load from gpt2 weight. w_b1 and w_b2 are the params of two linear layer
        Output: Tensor
    """
    w_b1, w_b2 = mlp['c_fc'], mlp['c_proj']
    
    out = linear(x, w_b1)
    out = gelu(out)
    out = linear(out, w_b2)

    return out


def attention(q, k, v, mask):  # [n_q, d_k], [n_k, d_k], [n_k, d_v], [n_q, n_k] -> [n_q, d_v]
    """
        Task: use torch API to implement attention computation according to formula(1) of the following paper
              where d_k account for the last dimension of `k`
        Paper: https://arxiv.org/abs/1706.03762
        Input: 
            q: Tensor
            k: Tensor
            v: Tensor
            mask: Tensor
            mlp: dictionary that load from gpt2 weight. w_b1 and w_b2 are the params of two linear layer
        Output: Tensor
    """
    #inner
    d_k = q.size(-1)
    scores = torch.matmul(q, k.t())
    scores = scores / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(torch.isinf(mask), float('-inf'))

    attn_weight = softmax(scores)
    out = torch.matmul(attn_weight, v)

    return out

def mha(x, attn, n_head):  # [n_seq, n_embd] -> [n_seq, n_embd]
    """
        Task: Complete the code of the multi-head attention
        
        Input: 
            x: Tensor
            attn: dictionary that load from gpt2 weight. c_attn and c_proj are the params of two linear layer
            n_head: number of head
        Output: Tensorying multi-head attention and linear transformation, shape [n_seq, n_embd].
    """
    c_attn, c_proj = attn['c_attn'], attn['c_proj']
    # qkv projection
    x = linear(x, c_attn)  # [n_seq, n_embd] -> [n_seq, 3*n_embd]
    
    # Split into qkv
    """
        Task: Split the q,k,v matrix from the tensor x
        Notes: [n_seq, 3*n_embd] -> 3 * [n_seq, n_embd]
    """
    q, k, v = x.split(x.size(-1) // 3, dim=-1)
    qkv = [q, k, v] # need to modify

    # Split into heads
    qkv_heads = [qkv_part.chunk(n_head, dim=-1) for qkv_part in qkv]  # 3 * [n_seq, n_embd] -> 3 * n_head * [n_seq, n_embd/n_head]
    qkv_heads = list(zip(*qkv_heads))  # [3, n_head, n_seq, n_embd/n_head]

    # Causal mask to hide future inputs from being attended to
    """
        Task: Construct mask matrix
        Notes: 
            | 0  -inf -inf ... -inf |
            | 0    0  -inf ... -inf |
            | 0    0    0  ... -inf |
            |...  ...  ... ...  ... | 
            | 0    0    0  ...   0  |
        Mask is a tensor whose dimension is [n_seq, n_seq]
    """
    n_seq = x.size(0)
    causal_mask = torch.zeros((n_seq, n_seq), device=x.device)
    causal_mask = causal_mask.masked_fill(
        torch.triu(torch.ones(n_seq, n_seq, dtype=torch.bool, device=x.device), diagonal=1),
        float('-inf')
    )

    # Perform attention over each head
    out_heads = [attention(q, k, v, causal_mask) for q, k, v in qkv_heads]  # n_head * [n_seq, n_embd/n_head]
    
    # Merge heads
    """
        Task: merge multi-heads results
        Notes: n_head * [n_seq, n_embd/n_head] --> [n_seq, n_embd]
    """
    x = torch.cat(out_heads, dim=-1)
    
    # Out projection
    x = linear(x, c_proj)  # [n_seq, n_embd] -> [n_seq, n_embd]
    
    return x


def transformer_block(x, block, n_head):  # [n_seq, n_embd] -> [n_seq, n_embd]
    mlp, attn, ln_1, ln_2 = block['mlp'], block['attn'], block['ln_1'], block['ln_2']
    
    # multi-head causal self attention
    x = x + mha(layer_norm(x, ln_1), attn, n_head=n_head)  # [n_seq, n_embd] -> [n_seq, n_embd]

    # position-wise feed forward network
    x = x + ffn(layer_norm(x, ln_2), mlp)  # [n_seq, n_embd] -> [n_seq, n_embd]

    return x


def gpt2(inputs, params, n_head):  # [n_seq] -> [n_seq, n_vocab]
    wte, wpe, blocks, ln_f = params['wte'], params['wpe'], params['blocks'], params['ln_f']
    # token + positional embeddings
    x = wte[inputs] + wpe[range(len(inputs))]  # [n_seq] -> [n_seq, n_embd]
    
    x = torch.Tensor(x)
    # forward pass through n_layer transformer blocks
    for block in blocks:
        x = transformer_block(x, block, n_head=n_head)  # [n_seq, n_embd] -> [n_seq, n_embd]

    # projection to vocab
    x = layer_norm(x, ln_f)  # [n_seq, n_embd] -> [n_seq, n_embd]
    return x @ wte.T  # [n_seq, n_embd] -> [n_seq, n_vocab]


def generate(inputs, params, n_head, n_tokens_to_generate):
    from tqdm import tqdm

    for _ in tqdm(range(n_tokens_to_generate), "generating"):  # auto-regressive decode loop
        logits = gpt2(inputs, params, n_head=n_head)  # model forward pass
        next_id = np.argmax(logits[-1])  # greedy sampling
        inputs.append(int(next_id))  # append prediction to input

    return inputs[len(inputs) - n_tokens_to_generate :]  # only return generated ids

def greedy_speculative_generate(inputs, draft_params, target_params, hparams_draft, hparams_target, n_tokens_to_generate, K):
    """
    Implements greedy speculative decoding using a draft model and a target model
    
    Args:
        inputs (list): Initial token IDs from the prompt
        draft_params: Parameters for the draft model (124M)
        target_params: Parameters for the target model (1558M)
        hparams_draft: Hyperparameters for draft model
        hparams_target: Hyperparameters for target model
        n_tokens_to_generate: Number of new tokens to generate
        K: Number of tokens to speculate at each step
    
    Returns:
       list: Generated token IDs
    """
    generated_ids = []
    current_inputs = list(inputs)
    
    # use tqdm to show where it is
    from tqdm import tqdm
    pbar = tqdm(total=n_tokens_to_generate, desc="greedy_speculative_generate model")
    
    while len(generated_ids) < n_tokens_to_generate:

        remaining = n_tokens_to_generate - len(generated_ids)
        if remaining <= 0:
            break

        # use124M
        draft_tokens = []
        draft_seq = current_inputs.copy()
        for _ in range(min(K, remaining)):
            logits = gpt2(draft_seq, draft_params, hparams_draft["n_head"])
            next_token = torch.argmax(logits[-1]).item()
            draft_tokens.append(next_token)
            draft_seq.append(next_token)
        
        # test it 
        full_seq = current_inputs + draft_tokens
        target_logits = gpt2(full_seq, target_params, hparams_target["n_head"])
        
        # judge
        new_tokens = []
        for i in range(len(draft_tokens)):
            pos = len(current_inputs) + i
            target_token = torch.argmax(target_logits[pos]).item()
            new_tokens.append(target_token)
            
            if target_token != draft_tokens[i]:
                break
        
        if len(new_tokens) == len(draft_tokens) and len(generated_ids) + len(new_tokens) < n_tokens_to_generate:
            extra_seq = current_inputs + new_tokens
            extra_logits = gpt2(extra_seq, target_params, hparams_target["n_head"])
            extra_token = torch.argmax(extra_logits[-1]).item()
            new_tokens.append(extra_token)
        
        # update
        generated_ids.extend(new_tokens)
        current_inputs.extend(new_tokens)
        pbar.update(len(new_tokens))
    
    pbar.close()

    return generated_ids



def main(prompt: str,
         n_tokens_to_generate: int = 5,
         model_size: str = "124M",
         enhance_model_size: str = None,
         models_dir: str = "models",
         K: int = 4):
    from utils import load_encoder_hparams_and_params

    # load encoder, hparams, and params from the released open-ai gpt-2 files
    encoder, hparams, params = load_encoder_hparams_and_params(model_size, models_dir)
    enhance_params = None
    if enhance_model_size is not None:
        _, enhance_hparams, enhance_params = load_encoder_hparams_and_params(
            enhance_model_size, 
            models_dir
        )

    # encode the input string using the BPE tokenizer
    input_ids = encoder.encode(prompt)

    # make sure we are not surpassing the max sequence length of our model
    assert len(input_ids) + n_tokens_to_generate < hparams["n_ctx"]

    # generate output ids
    start = time.time()
    if enhance_model_size is None or enhance_params is None:
        output_ids = generate(input_ids, params, hparams["n_head"], n_tokens_to_generate)
    else:
        output_ids = greedy_speculative_generate(input_ids, params, enhance_params, hparams, enhance_hparams, n_tokens_to_generate, K )
    end = time.time()
    print(f"Time taken to generate {n_tokens_to_generate} tokens: {end - start:.2f}s")

    # decode the ids back into a string
    output_text = encoder.decode(output_ids)
    return output_text


if __name__ == "__main__":
    import fire
    fire.Fire(main)
